#include <stdio.h>
#include <stdlib.h>
#include <string.h>

// Simple token types
typedef enum {
    TOKEN_INT, TOKEN_CHAR, TOKEN_FLOAT, TOKEN_DOUBLE, TOKEN_VOID,
    TOKEN_IF, TOKEN_ELSE, TOKEN_WHILE, TOKEN_FOR, TOKEN_RETURN,
    TOKEN_PLUS, TOKEN_MINUS, TOKEN_MULTIPLY, TOKEN_DIVIDE,
    TOKEN_ASSIGN, TOKEN_EQUAL, TOKEN_SEMICOLON, TOKEN_LEFT_PAREN,
    TOKEN_RIGHT_PAREN, TOKEN_LEFT_BRACE, TOKEN_RIGHT_BRACE,
    TOKEN_INTEGER_LITERAL, TOKEN_IDENTIFIER, TOKEN_EOF
} TokenType;

// Token structure
typedef struct {
    TokenType type;
    char value[256];
    int line;
    int column;
} Token;

// Simple lexer
Token* tokenize(const char* source) {
    static Token tokens[1000];
    int tokenCount = 0;
    int pos = 0;
    int line = 1;
    int col = 1;
    
    while (source[pos] != '\0') {
        // Skip whitespace
        while (source[pos] == ' ' || source[pos] == '\t' || source[pos] == '\n') {
            if (source[pos] == '\n') {
                line++;
                col = 1;
            } else {
                col++;
            }
            pos++;
        }
        
        if (source[pos] == '\0') break;
        
        // Keywords and identifiers
        if ((source[pos] >= 'a' && source[pos] <= 'z') || 
            (source[pos] >= 'A' && source[pos] <= 'Z') || 
            source[pos] == '_') {
            
            int start = pos;
            while ((source[pos] >= 'a' && source[pos] <= 'z') || 
                   (source[pos] >= 'A' && source[pos] <= 'Z') || 
                   (source[pos] >= '0' && source[pos] <= '9') || 
                   source[pos] == '_') {
                pos++;
            }
            
            int len = pos - start;
            char word[256];
            strncpy(word, source + start, len);
            word[len] = '\0';
            
            tokens[tokenCount].line = line;
            tokens[tokenCount].column = col;
            strcpy(tokens[tokenCount].value, word);
            
            if (strcmp(word, "int") == 0) tokens[tokenCount].type = TOKEN_INT;
            else if (strcmp(word, "char") == 0) tokens[tokenCount].type = TOKEN_CHAR;
            else if (strcmp(word, "float") == 0) tokens[tokenCount].type = TOKEN_FLOAT;
            else if (strcmp(word, "double") == 0) tokens[tokenCount].type = TOKEN_DOUBLE;
            else if (strcmp(word, "void") == 0) tokens[tokenCount].type = TOKEN_VOID;
            else if (strcmp(word, "if") == 0) tokens[tokenCount].type = TOKEN_IF;
            else if (strcmp(word, "else") == 0) tokens[tokenCount].type = TOKEN_ELSE;
            else if (strcmp(word, "while") == 0) tokens[tokenCount].type = TOKEN_WHILE;
            else if (strcmp(word, "for") == 0) tokens[tokenCount].type = TOKEN_FOR;
            else if (strcmp(word, "return") == 0) tokens[tokenCount].type = TOKEN_RETURN;
            else tokens[tokenCount].type = TOKEN_IDENTIFIER;
            
            tokenCount++;
            col += len;
        }
        // Numbers
        else if (source[pos] >= '0' && source[pos] <= '9') {
            int start = pos;
            while (source[pos] >= '0' && source[pos] <= '9') {
                pos++;
            }
            
            int len = pos - start;
            tokens[tokenCount].type = TOKEN_INTEGER_LITERAL;
            tokens[tokenCount].line = line;
            tokens[tokenCount].column = col;
            strncpy(tokens[tokenCount].value, source + start, len);
            tokens[tokenCount].value[len] = '\0';
            
            tokenCount++;
            col += len;
        }
        // Operators and symbols
        else {
            tokens[tokenCount].line = line;
            tokens[tokenCount].column = col;
            tokens[tokenCount].value[0] = source[pos];
            tokens[tokenCount].value[1] = '\0';
            
            switch (source[pos]) {
                case '+': tokens[tokenCount].type = TOKEN_PLUS; break;
                case '-': tokens[tokenCount].type = TOKEN_MINUS; break;
                case '*': tokens[tokenCount].type = TOKEN_MULTIPLY; break;
                case '/': tokens[tokenCount].type = TOKEN_DIVIDE; break;
                case '=': tokens[tokenCount].type = TOKEN_ASSIGN; break;
                case ';': tokens[tokenCount].type = TOKEN_SEMICOLON; break;
                case '(': tokens[tokenCount].type = TOKEN_LEFT_PAREN; break;
                case ')': tokens[tokenCount].type = TOKEN_RIGHT_PAREN; break;
                case '{': tokens[tokenCount].type = TOKEN_LEFT_BRACE; break;
                case '}': tokens[tokenCount].type = TOKEN_RIGHT_BRACE; break;
                default: tokens[tokenCount].type = TOKEN_EOF; break;
            }
            
            tokenCount++;
            pos++;
            col++;
        }
    }
    
    // Add EOF token
    tokens[tokenCount].type = TOKEN_EOF;
    tokens[tokenCount].value[0] = '\0';
    tokens[tokenCount].line = line;
    tokens[tokenCount].column = col;
    
    return tokens;
}

// Simple code generator
void generateAssembly(const Token* tokens) {
    printf(".section .text\n");
    printf(".global _start\n\n");
    printf("main:\n");
    printf("    push %%rbp\n");
    printf("    mov %%rsp, %%rbp\n");
    
    int i = 0;
    while (tokens[i].type != TOKEN_EOF) {
        if (tokens[i].type == TOKEN_INT && tokens[i+1].type == TOKEN_IDENTIFIER) {
            // Variable declaration
            printf("    ; int %s\n", tokens[i+1].value);
            if (tokens[i+2].type == TOKEN_ASSIGN && tokens[i+3].type == TOKEN_INTEGER_LITERAL) {
                printf("    mov $%s, %%rax\n", tokens[i+3].value);
                printf("    mov %%rax, -8(%%rbp)\n");
                i += 4;
            } else {
                i += 2;
            }
        } else if (tokens[i].type == TOKEN_RETURN) {
            printf("    ; return\n");
            if (tokens[i+1].type == TOKEN_INTEGER_LITERAL) {
                printf("    mov $%s, %%rax\n", tokens[i+1].value);
                i += 2;
            } else if (tokens[i+1].type == TOKEN_IDENTIFIER) {
                printf("    mov -8(%%rbp), %%rax\n");
                i += 2;
            } else {
                i++;
            }
            printf("    mov %%rbp, %%rsp\n");
            printf("    pop %%rbp\n");
            printf("    ret\n");
        } else {
            i++;
        }
    }
    
    printf("\n_start:\n");
    printf("    mov $60, %%rax\n");
    printf("    mov $0, %%rdi\n");
    printf("    syscall\n");
}

int main(int argc, char* argv[]) {
    if (argc < 2) {
        printf("Usage: %s <input_file> [output_file]\n", argv[0]);
        return 1;
    }
    
    // Read input file
    FILE* file = fopen(argv[1], "r");
    if (!file) {
        printf("Error: Could not open file %s\n", argv[1]);
        return 1;
    }
    
    // Get file size
    fseek(file, 0, SEEK_END);
    long size = ftell(file);
    fseek(file, 0, SEEK_SET);
    
    // Read file content
    char* source = malloc(size + 1);
    fread(source, 1, size, file);
    source[size] = '\0';
    fclose(file);
    
    printf("Source code:\n%s\n\n", source);
    
    // Tokenize
    printf("Tokenizing...\n");
    Token* tokens = tokenize(source);
    
    // Print tokens
    int i = 0;
    while (tokens[i].type != TOKEN_EOF) {
        printf("Token: %d, Value: %s, Line: %d, Col: %d\n", 
               tokens[i].type, tokens[i].value, tokens[i].line, tokens[i].column);
        i++;
    }
    printf("Token: %d, Value: %s, Line: %d, Col: %d\n", 
           tokens[i].type, tokens[i].value, tokens[i].line, tokens[i].column);
    
    printf("\nGenerating assembly...\n");
    generateAssembly(tokens);
    
    free(source);
    return 0;
}
